import pyaudio
import wave
import datetime
import numpy as np
import matplotlib.pyplot as plt

# 配置音频输入
format = pyaudio.paInt16
channels = 1
rate = 44100
chunk_size = 1024
record_seconds = 1

# 创建一个PyAudio对象
audio = pyaudio.PyAudio()

# 打开音频输入流
stream = audio.open(format=format,
                    channels=channels,
                    rate=rate,
                    input=True,
                    frames_per_buffer=chunk_size)

# 读取指定秒数的音频数据
frames = []
for i in range(int(rate / chunk_size * record_seconds)):
    data = stream.read(chunk_size)
    frames.append(data)

# 关闭音频流和PyAudio对象
stream.stop_stream()
stream.close()
audio.terminate()

# 将音频数据保存为WAV文件
current_time = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
filename = f"audio_{current_time}.wav"

with wave.open(filename, "wb") as wf:
    wf.setnchannels(channels)
    wf.setsampwidth(audio.get_sample_size(format))
    wf.setframerate(rate)
    wf.writeframes(b''.join(frames))

print(f"音频已保存为 {filename}")

# 读取音频文件数据
with wave.open(filename, 'rb') as wf:
    # 获取音频参数
    channels = wf.getnchannels()
    sample_width = wf.getsampwidth()
    frame_rate = wf.getframerate()
    num_frames = wf.getnframes()

    # 读取音频数据
    audio_data = wf.readframes(num_frames)

# 将音频数据转换为一维numpy数组
if sample_width == 1:
    audio_data = np.frombuffer(audio_data, dtype=np.uint8)
elif sample_width == 2:
    audio_data = np.frombuffer(audio_data, dtype=np.int16)
elif sample_width == 4:
    audio_data = np.frombuffer(audio_data, dtype=np.int32)

# 创建时间轴
duration = num_frames / frame_rate
time = np.linspace(0., duration, num=len(audio_data))

# 创建波形图
fig, ax = plt.subplots()
ax.plot(time, audio_data)

# 设置图形相关参数
ax.set_xlabel('Time (s)')
ax.set_ylabel('Amplitude')
ax.set_title('Audio Waveform')

# 显示波形图
plt.show()
